Skip to content

[PyTorch] NCCL EP zero copy with symmem pool and user provided recv_topk_weight tensor#3187

Merged
phu0ngng merged 10 commits into
NVIDIA:mainfrom
YangFei1990:nccl_ep_zero_cp
Jul 21, 2026
Merged

[PyTorch] NCCL EP zero copy with symmem pool and user provided recv_topk_weight tensor#3187
phu0ngng merged 10 commits into
NVIDIA:mainfrom
YangFei1990:nccl_ep_zero_cp

Conversation

@YangFei1990

@YangFei1990 YangFei1990 commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Description

Addition support for NCCL EP zero copy feature, including

  • Expose user provided recv_topk_weights buffer
  • Use symm memory pool to allocate symm buffer

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@YangFei1990
expose user provided weights
1602fc8 
Signed-off-by: YangFei1990 <feiw@nvidia.com>
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jul 7, 2026
@YangFei1990 YangFei1990 mentioned this pull request Jul 9, 2026
Open
6 tasks
@YangFei1990
add check for symm backed tensor
0c24ee5 
Signed-off-by: YangFei1990 <feiw@nvidia.com>
@YangFei1990
Merge branch 'nccl_ep_zero_cp' of https://github.com/YangFei1990/Tran…
c37c460 
…sformerEngine into nccl_ep_zero_cp
@YangFei1990
add zero copy tests
e12d106 
Signed-off-by: YangFei1990 <feiw@nvidia.com>
@YangFei1990
YangFei1990 marked this pull request as ready for review July 15, 2026 21:32
@YangFei1990
YangFei1990 requested a review from ksivaman as a code owner July 15, 2026 21:32
@greptile-apps

greptile-apps Bot commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR replaces the static, EpBuffer-owned symm-mem buffer strategy with on-demand allocation from a process-global symm-mem MemPool, and exposes caller-supplied recv_tokens, recv_topk_weights, and grad_out kwargs on ep_dispatch/ep_combine so mcore can manage its own IO buffers.

  • Pool allocation (_get_symm_mem_pool + _alloc_io): A single process-wide MemPool backed by the NCCL symm-mem allocator is created once at first dispatch; tensors are allocated per-call, ref-counted by PyTorch, and auto-registered with the NCCL comm — removing the need for per-buffer rendezvous and static lifetime management in EpBuffer.
  • API simplification: EpBuffer.__init__ drops dispatch_recv_tokens / combine_grad_expert_out kwargs and all three *_symm_buf slots; callers pass buffers directly at ep_dispatch/ep_combine call time, and the new is_symm_backed helper lets callers verify pool-backed tensors.
  • Tests updated: Tests migrate from buffer-construction kwargs to call-site kwargs; new test_zero_copy_pool_auto_alloc covers the mcore-primary path where neither recv nor grad buffer is supplied.

Confidence Score: 5/5

Safe to merge. The refactor correctly moves symm-mem buffer ownership out of EpBuffer and into on-demand pool allocation, with no changes to the C++ dispatch/combine kernels.

The allocation logic is consistent across forward and backward, mark_not_offload is applied everywhere pool tensors are created, and the backend-mismatch guard added in the previous review round is in place. No correctness regression was found in the dispatch or combine paths.

No files require special attention. The two observations are minor: ep_group is unused in the pool path of symm_mem_alloc, and the pool auto-alloc test could add a symmetry check for recv_topk_weights.

Important Files Changed

Filename Overview
transformer_engine/pytorch/distributed.py Adds _get_symm_mem_pool (process-global, created once with backend assertion) and extends symm_mem_alloc with use_pool/backend kwargs. Pool creation is defensive (tries two API shapes). ep_group is silently dropped in the pool path — a minor API design issue but no runtime impact on the current call sites.
transformer_engine/pytorch/ep.py Replaces static EpBuffer-owned symm-mem buffers with on-demand pool allocation via _alloc_io. ep_dispatch/ep_combine now accept optional caller-supplied recv_tokens, recv_topk_weights, and grad_out kwargs. is_symm_backed added to __all__. Logic is sound; pool allocations call mark_not_offload correctly.
tests/pytorch/distributed/run_ep.py Tests migrated from buffer-construction kwargs to dispatch/combine call-site kwargs, matching the new API. New test_zero_copy_pool_auto_alloc covers the primary mcore path. recv_w symmetry not asserted in the pool auto-alloc test.

Reviews (5): Last reviewed commit: "Merge branch 'main' into nccl_ep_zero_cp" | Re-trigger Greptile

greptile-apps[bot]
greptile-apps Bot reviewed Jul 15, 2026
View reviewed changes
Comment thread transformer_engine/pytorch/distributed.py Outdated
Comment thread transformer_engine/pytorch/ep.py
Comment thread tests/pytorch/distributed/run_ep.py
@phu0ngng
phu0ngng self-requested a review July 20, 2026 15:03
phu0ngng
phu0ngng previously approved these changes Jul 21, 2026
View reviewed changes
@phu0ngng

phu0ngng commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

/te-ci pytorch L1

phu0ngng and others added 2 commits July 21, 2026 10:21
@phu0ngng
Merge branch 'main' into nccl_ep_zero_cp
c719c59
@YangFei1990
Address review: guard symm-mem pool backend + is_symm_backed bootstrap
2f80d88 
- _get_symm_mem_pool: raise if a later call requests a different backend than
  the one the process-wide pool was created with (was silently ignored).
- is_symm_backed: raise a clear error if called before bootstrap (_EP_GROUP is
  None) instead of letting the bare except swallow the AttributeError.

Signed-off-by: YangFei1990 <feiw@nvidia.com>
@phu0ngng

phu0ngng commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

/te-ci pytorch L1

@phu0ngng phu0ngng mentioned this pull request Jul 21, 2026
Open
13 tasks
@phu0ngng phu0ngng changed the title NCCL EP zero copy additional support [PyTorch] NCCL EP zero copy additional support Jul 21, 2026
@phu0ngng phu0ngng changed the title [PyTorch] NCCL EP zero copy additional support [PyTorch] NCCL EP zero copy with symmem pool Jul 21, 2026
@phu0ngng phu0ngng changed the title [PyTorch] NCCL EP zero copy with symmem pool [PyTorch] NCCL EP zero copy with symmem pool and user provided recv_topk_weight tensor Jul 21, 2026
@phu0ngng
phu0ngng merged commit 622a3ee into NVIDIA:main Jul 21, 2026
10 of 14 checks passed
@KshitijLakhani

KshitijLakhani commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

@phu0ngng and @YangFei1990 is this supposed to be cherry-picked for TE 2.18
If yes, please tag it. Thank you :)

cc: @fheinecke (Just FYI - no action needed)

@phu0ngng phu0ngng added the 2.18 label Jul 21, 2026
KshitijLakhani pushed a commit that referenced this pull request Jul 23, 2026
@YangFei1990 @phu0ngng @KshitijLakhani
[PyTorch] NCCL EP zero copy with symmem pool and user provided recv_t…
81453e2 
…opk_weight tensor (#3187)

* expose user-provided weights

* adding pool based symm allocation; remove the persistent buffer in EpBuffer

* add zero copy tests

Signed-off-by: YangFei1990 <feiw@nvidia.com>

---------

Signed-off-by: YangFei1990 <feiw@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@greptile-apps greptile-apps[bot] greptile-apps[bot] left review comments
@phu0ngng phu0ngng phu0ngng approved these changes
@ksivaman ksivaman Awaiting requested review from ksivaman ksivaman is a code owner
+1 more reviewer
@Autumn1998 Autumn1998 Autumn1998 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

2.18 community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@YangFei1990 @phu0ngng @KshitijLakhani @Autumn1998